#ifndef CAMERA_CONFIG_HPP
#define CAMERA_CONFIG_HPP

#include "utility.hpp"
#include "parallel_utility.hpp"
#include "math.h"

namespace LfCamStruct
{

    struct Camera_config
    {
        cv::Mat K_;         // uncalibrated intrinsic matrix
        cv::Mat Optimal_K_; // calibrated intrinsic matrix
        cv::Mat distortion_coeff_;
        bool distortion_flag_;
        float camera_height_; // m
        float yaw_;           // deg
        float pitch_;         // deg
        float roll_;          // deg
        float tan_alpha_v_;
        float tan_alpha_h_;
        float FOV_v_; // deg
        float FOV_h_; // deg
        cv::Size sz_;

        // undistortion related
        cv::Mat undistort_map1_;
        cv::Mat undistort_map2_;
        int stripe_size0_;

        void set_FOV()
        {
            float &cx = Optimal_K_.at<float>(0, 2);
            float &cy = Optimal_K_.at<float>(1, 2);
            float &fx = Optimal_K_.at<float>(0, 0);
            float &fy = Optimal_K_.at<float>(1, 1);
            tan_alpha_v_ = cy / fy;
            tan_alpha_h_ = cx / fx;
            FOV_h_ = 2 * atan(tan_alpha_h()) * LfUtils::DegperRad;
            FOV_v_ = 2 * atan(tan_alpha_v()) * LfUtils::DegperRad;
        }

    public:
        Camera_config() {}
        Camera_config(const cv::Mat &K,
                      const cv::Mat &Distortion_coeff_,
                      float Cam_height,
                      float Yaw,
                      float Pitch,
                      float Roll,
                      const cv::Size &Sz) : K_(K), distortion_coeff_(Distortion_coeff_), camera_height_(std::abs(Cam_height)),
                                            yaw_(Yaw), pitch_(Pitch), roll_(Roll), sz_(Sz)
        {
            //distortion_flag_ = true;
            //Optimal_K_ = cv::getOptimalNewCameraMatrix(K_, distortion_coeff_, sz_, 0);
            distortion_flag_ = false;
            Optimal_K_  = K_;
            set_FOV();
            undistort_map_init();
        }

        const float &cam_height() const { return camera_height_; }
        float &cam_height() { return camera_height_; }
        const float &yaw() const { return yaw_; }
        float &yaw() { return yaw_; }
        const float &pitch() const { return pitch_; }
        float &pitch() { return pitch_; }
        const float &roll() const { return roll_; }
        float &roll() { return roll_; }
        const float &tan_alpha_v() const { return tan_alpha_v_; }
        const float &tan_alpha_h() const { return tan_alpha_h_; }
        const float &FOV_v() const { return FOV_v_; }
        const float &FOV_h() const { return FOV_h_; }
        const cv::Size &sz() const { return sz_; }
        const cv::Mat &K() const { return Optimal_K_; }

        bool set_distortion(bool calibrated_flag)
        {
            if (distortion_flag_ != calibrated_flag)
            {
                distortion_flag_ = calibrated_flag;
                if (distortion_flag_)
                {
                    Optimal_K_ = cv::getOptimalNewCameraMatrix(K_, distortion_coeff_, sz_, true);
                }
                else
                {
                    Optimal_K_ = K_;
                }
                set_FOV();
            }
            return distortion_flag_;
        }

        std::vector<cv::Point2f> undistortPts(const std::vector<cv::Point2f> &pts) const
        {
            if (distortion_flag_ && !pts.empty())
            {
                std::vector<cv::Point2f> undistorted_pts;
                cv::undistortPoints(pts, undistorted_pts, K_, distortion_coeff_, cv::noArray(), Optimal_K_);
                return undistorted_pts;
            }
            return pts;
        }

        cv::Rect2f undistortRect(const cv::Rect2f &rect)
        {
            if (distortion_flag_)
            {
                std::vector<cv::Point2f> tl_br{rect.tl(), rect.br()};
                std::vector<cv::Point2f> undistored_tl_br;
                cv::undistortPoints(tl_br, undistored_tl_br, K_, distortion_coeff_, cv::noArray(), Optimal_K_);
                return cv::Rect2f(undistored_tl_br[0], undistored_tl_br[1]);
            }
            return rect;
        }

        cv::Mat undistortImg(const cv::Mat &img)
        {
            if (distortion_flag_)
            {
                cv::Mat rectify_img = cv::Mat(img.size(), img.type());

                cv::remap(img, rectify_img, undistort_map1_, undistort_map2_, cv::INTER_LINEAR, cv::BORDER_CONSTANT);

                // for (int y = 0; y < sz_.height; y += stripe_size0_) {
                //     int stripe_size = std::min( stripe_size0_, sz_.height - y );
                //     cv::Mat map1_part = undistort_map1_.rowRange(y, y + stripe_size),
                //             map2_part = undistort_map2_.rowRange(y, y + stripe_size),
                //             dst_part  = rectify_img.rowRange(y, y + stripe_size);
                //     cv::remap(img, dst_part, map1_part, map2_part, cv::INTER_LINEAR, cv::BORDER_CONSTANT);
                // }

                // auto do_work = [this, &rectify_img, &img](int i) {
                //     int y = i * stripe_size0_;
                //     int stripe_size = std::min( stripe_size0_, sz_.height - y );
                //     cv::Mat map1_part = undistort_map1_.rowRange(y, y + stripe_size),
                //             map2_part = undistort_map2_.rowRange(y, y + stripe_size),
                //             dst_part  = rectify_img.rowRange(y, y + stripe_size);
                //     cv::remap(img, dst_part, map1_part, map2_part, cv::INTER_LINEAR, cv::BORDER_CONSTANT) ;
                // };

                // /* do_work parallelly */
                // size_t total_num = (sz_.height / stripe_size0_);
                // size_t chunk = LOOP_PARALISM::get_chunk(total_num);
                // for (size_t i = 0; i < chunk; ++i) {
                //     LOOP_PARALISM loop_paraller;
                //     for (size_t j = 0; j < loop_paraller.numthreads_; ++j) {
                //         size_t index = i * loop_paraller.numthreads_ + j;
                //         if (index >= total_num) { continue; }
                //         loop_paraller.works_[j] = std::thread(do_work, index);
                //     }
                // }
                return rectify_img;
            }
            return img;
        }

    private:
        void undistort_map_init()
        {
            cv::Mat cameraMatrix = K_;
            cv::Mat distCoeffs = distortion_coeff_;
            cv::Mat newCameraMatrix = Optimal_K_;

            stripe_size0_ = std::min(std::max(1, (1 << 12) / std::max(sz_.width, 1)), sz_.height);
            undistort_map1_ = cv::Mat(sz_.height, sz_.width, CV_16SC2);
            undistort_map2_ = cv::Mat(sz_.height, sz_.width, CV_16UC1);

            cv::Mat_<double> A, Ar, I = cv::Mat_<double>::eye(3, 3);
            cameraMatrix.convertTo(A, CV_64F);
            if (!distCoeffs.empty())
                distCoeffs = cv::Mat_<double>(distCoeffs);
            else
            {
                distCoeffs.create(5, 1, CV_64F);
                distCoeffs = 0.;
            }

            if (!newCameraMatrix.empty())
                newCameraMatrix.convertTo(Ar, CV_64F);
            else
                A.copyTo(Ar);

            double v0 = Ar(1, 2);
            for (int y = 0; y < sz_.height; y += stripe_size0_)
            {
                int stripe_size = std::min(stripe_size0_, sz_.height - y);
                Ar(1, 2) = v0 - y;
                cv::Mat map1_part = undistort_map1_.rowRange(y, y + stripe_size),
                        map2_part = undistort_map2_.rowRange(y, y + stripe_size);
                cv::initUndistortRectifyMap(A, distCoeffs, I, Ar, cv::Size(sz_.width, stripe_size),
                                            map1_part.type(), map1_part, map2_part);
            }
        }
    };

    inline std::ostream &operator<<(std::ostream &os, const Camera_config &data)
    {
        os << "================ Camera_config Info ================" << std::endl
           << "K: \n"
           << data.K_ << std::endl
           << "Optimal_K: \n"
           << data.Optimal_K_ << std::endl
           << "distortion_coeff: \n"
           << data.distortion_coeff_ << std::endl
           << "camera_height: " << data.camera_height_ << " m" << std::endl
           << "yaw: " << data.yaw_ << " deg" << std::endl
           << "pitch: " << data.pitch_ << " deg" << std::endl
           << "roll: " << data.roll_ << " deg" << std::endl
           << "FOV_v; " << data.FOV_v_ << " deg " << std::endl
           << "FOV_h; " << data.FOV_h_ << " deg " << std::endl
           << "sz: " << data.sz_ << std::endl;
        return os;
    }
} // namespace LfCamStruct
#endif // CAMERA_CONFIG_HPP
